package com.bleeth.flow.step.field;

import cn.hutool.core.util.StrUtil;
import com.bleeth.flow.core.common.APlugin;
import com.bleeth.flow.core.common.PluginAnnotation;
import com.bleeth.flow.core.common.PluginTypeEnum;
import com.jd.platform.async.wrapper.WorkerWrapper;
import lombok.Data;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;

import java.io.Serializable;
import java.util.List;
import java.util.Map;

/**
 * @author ：Bleeth
 * @date ：2021-08-05
 * @description：字段选择，删除组件
 */
@Data
@PluginAnnotation(name = "自定义聚合",
        type = PluginTypeEnum.AGG,
        description = "",
        id = "SelectFieldPlugin")
public class SelectFieldPlugin extends APlugin implements Serializable {

    private static final long serialVersionUID = -5278879206293829464L;

    private List<SelectFieldParamBean> fieldParamList;


    @Override
    public Dataset<Row> action(Dataset<Row> input, Map<String, WorkerWrapper> allWrappers) {
        super.action(input, allWrappers);

        String fromPluginName = fromList.get(0);
        WorkerWrapper fromWrapper = allWrappers.get(fromPluginName);
        Dataset<Row> ds = (Dataset<Row>) fromWrapper.getWorkResult().getResult();


        for (int i = 0; i < fieldParamList.size(); i++) {
            SelectFieldParamBean fieldParamBean = fieldParamList.get(i);
            String fieldName = fieldParamBean.getFieldName();
            String fieldType = fieldParamBean.getFieldType();
            String aliasName = fieldParamBean.getAliasName();

            if (StrUtil.equalsIgnoreCase(fieldType, "REMOVE")) {
                ds = ds.drop(fieldName);
            }

            if (StrUtil.equalsIgnoreCase(fieldType, "CHANGE")) {
                ds = ds.withColumnRenamed(fieldName,aliasName);
            }
        }
        return ds;
    }


    @Data
    public static class SelectFieldParamBean implements Serializable {
        private static final long serialVersionUID = 543195102494187159L;

        //字段名称
        private String fieldName;

        //操作类型 REMOVE(移除该字段)  CHANGE(修改字段名称)
        private String fieldType;

        //修改字段名称
        private String aliasName;
    }





}
